import matplotlib.pyplot as plt
import pickle
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

# 加载手写数字数据集
digits = load_digits()
X, y = digits.data, digits.target

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 初始化一个列表来存储每个 K 值的准确率
accuracies = []

# 尝试不同的 K 值
for k in range(1, 41):
    knn = KNeighborsClassifier(n_neighbors=k)
    knn.fit(X_train, y_train)
    y_pred = knn.predict(X_test)
    accuracy = (y_pred == y_test).mean()
    accuracies.append(accuracy)
# 绘制准确率曲线
plt.figure(figsize=(10, 6))
plt.plot(range(1, 41), accuracies, marker='o')
plt.title('KNN Accuracy for Different K Values')
plt.xlabel('K Value')
plt.ylabel('Accuracy')
plt.grid(True)
best_k = range(1, 41)[accuracies.index(max(accuracies))]
plt.axvline(x=best_k, color='black', linestyle='--')

plt.savefig('accuracy_plot.pdf', format='pdf')

# 训练最优的 KNN 模型
best_knn = KNeighborsClassifier(n_neighbors=best_k)
best_knn.fit(X_train, y_train)

# 保存最优的 KNN 模型
with open('best_knn_model.pickle', 'wb') as f:
    pickle.dump(best_knn, f)

print(f"最优的 K 值是 {best_k}，准确率为 {max(accuracies)}")